#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
MySQL批量多步骤操作脚本（5万条数据处理版）

功能：
1. 批量执行四个关联的SQL步骤
2. 高效处理5万条数据量
3. 详细日志记录和进度跟踪
4. 自动分批处理避免内存溢出
5. 事务管理和错误恢复机制
"""
import pandas as pd
import pymysql
import logging
from typing import List, Tuple, Optional
from datetime import datetime
import time
import sys
from math import ceil

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('mysql_batch_processing.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class MySQLBatchProcessor:
    def __init__(self, host: str, port: int, user: str, password: str, database: str):
        self.db_config = {
            'host': host,
            'port': port,
            'user': user,
            'password': password,
            'database': database,
            'charset': 'utf8mb4',
            'cursorclass': pymysql.cursors.DictCursor
        }
        self.connection = None
        self.batch_size = 2000  # 根据服务器性能调整
        self.query_timeout = 300  # 查询超时时间(秒)

    def connect(self):
        """建立数据库连接"""
        try:
            self.connection = pymysql.connect(**self.db_config)
            logger.info("成功连接到MySQL数据库")
        except pymysql.Error as e:
            logger.error(f"数据库连接失败: {str(e)}")
            raise

    def close(self):
        """关闭数据库连接"""
        if self.connection:
            try:
                self.connection.close()
                logger.info("数据库连接已关闭")
            except pymysql.Error as e:
                logger.error(f"关闭连接时出错: {str(e)}")
            finally:
                self.connection = None

    def execute_batch_update(self, sql: str, ids: List[str], description: str) -> int:
        """
        执行批量更新操作
        返回处理的行数
        """
        if not self.connection:
            self.connect()

        total = len(ids)
        processed = 0
        start_time = time.time()

        logger.info(f"开始执行 {description}，共 {total} 条记录")

        try:
            with self.connection.cursor() as cursor:
                # 设置超时
                cursor.execute(f"SET SESSION max_execution_time={self.query_timeout * 1000}")

                # 分批处理
                for i in range(0, total, self.batch_size):
                    batch = ids[i:i + self.batch_size]
                    placeholders = ','.join(['%s'] * len(batch))
                    formatted_sql = sql.replace('(?,?,?)', f'({placeholders})')

                    try:
                        cursor.execute(formatted_sql, batch)
                        self.connection.commit()
                        processed += len(batch)

                        # 打印进度
                        progress = processed / total * 100
                        if i % (10 * self.batch_size) == 0 or processed == total:
                            logger.info(
                                f"进度: {processed}/{total} ({progress:.1f}%) "
                                f"批次: {ceil(i / self.batch_size) + 1}/{ceil(total / self.batch_size)}"
                            )
                    except pymysql.Error as e:
                        self.connection.rollback()
                        logger.error(f"处理批次 {ceil(i / self.batch_size) + 1} 失败: {str(e)}")
                        raise

            elapsed = time.time() - start_time
            logger.info(f"{description} 完成，共处理 {processed} 条记录，耗时 {elapsed:.2f} 秒")
            return processed

        except Exception as e:
            logger.error(f"{description} 执行失败: {str(e)}")
            raise

    def process_step3(self, ord_item_xm_ids: List[str]) -> int:
        """
        处理步骤3：查询关联订单并更新
        返回更新的订单数
        """
        if not self.connection:
            self.connect()

        total = len(ord_item_xm_ids)
        processed = 0
        updated_orders = 0
        start_time = time.time()

        logger.info(f"开始执行步骤3，共 {total} 条ORD_ITM_XM_ID")

        try:
            with self.connection.cursor() as cursor:
                # 设置超时
                cursor.execute(f"SET SESSION max_execution_time={self.query_timeout * 1000}")

                # 存储所有找到的ORDER_ID
                order_ids = set()

                # 分批查询ORDER_ID
                for i in range(0, total, self.batch_size):
                    batch = ord_item_xm_ids[i:i + self.batch_size]
                    placeholders = ','.join(['%s'] * len(batch))

                    query_sql = f"""
                    SELECT co.ORDER_ID FROM 
                    CUSTOMER_ORDER_ITEM_XM coix
                    JOIN CUSTOMER_ORDER_ITEM coi ON coi.ORDER_ITEM_ID = coix.PAR_ORD_ITM_XM_ID 
                        AND coi.MEMBER_ID = coix.MEMBER_ID
                    JOIN CUSTOMER_ORDER co ON co.ORDER_ID = coi.ORDER_ID 
                        AND coi.MEMBER_ID = co.LOY_MEMBER_ID
                    WHERE coix.ORD_ITM_XM_ID IN ({placeholders})
                    """

                    try:
                        cursor.execute(query_sql, batch)
                        batch_order_ids = [row['ORDER_ID'] for row in cursor]
                        order_ids.update(batch_order_ids)
                        processed += len(batch)

                        # 打印进度
                        progress = processed / total * 100
                        if i % (10 * self.batch_size) == 0 or processed == total:
                            logger.info(
                                f"查询进度: {processed}/{total} ({progress:.1f}%) "
                                f"找到订单: {len(order_ids)}"
                            )
                    except pymysql.Error as e:
                        logger.error(f"查询批次 {ceil(i / self.batch_size) + 1} 失败: {str(e)}")
                        raise

                # 如果有找到订单则更新
                if order_ids:
                    order_ids = list(order_ids)
                    update_total = len(order_ids)
                    update_sql = """
                    UPDATE CUSTOMER_ORDER SET LAST_UPADTED = UTC_TIMESTAMP
                    WHERE ORDER_ID IN ({})
                    """.format(','.join(['%s'] * len(order_ids)))

                    try:
                        cursor.execute(update_sql, order_ids)
                        self.connection.commit()
                        updated_orders = cursor.rowcount
                        logger.info(f"更新了 {updated_orders} 条订单记录")
                    except pymysql.Error as e:
                        self.connection.rollback()
                        logger.error(f"订单更新失败: {str(e)}")
                        raise

            elapsed = time.time() - start_time
            logger.info(
                f"步骤3完成，处理 {processed} 条ORD_ITM_XM_ID，"
                f"更新 {updated_orders} 条订单，耗时 {elapsed:.2f} 秒"
            )
            return updated_orders

        except Exception as e:
            logger.error(f"步骤3执行失败: {str(e)}")
            raise


def load_ids_from_file(file_path: str) -> List[str]:
    # df = pd.read_csv(r"F:\workspace\pythonProject\hello\com\cn\for_cdc\input_data\KVB-MISSING_PROMO_ID.csv")
    # loy_rdm_itm_ids = df['ROW_ID'].values.tolist()

    loy_rdm_itm_ids=[]

    return loy_rdm_itm_ids


def main():
    # 数据库配置
    db_host = "10.128.56.153"
    db_port = 31410
    db_user = "cdcapp_kvb"
    db_password = "p438jUPYgI8SVD^t"
    db_name = "CDC_KVB_API"

    # 从文件加载ID（示例路径，请替换为实际路径）
    loy_rdm_itm_ids = load_ids_from_file()
    # loy_acrl_itm_ids = load_ids_from_file("loy_acrl_itm_ids.txt")[:50000]
    # ord_item_xm_ids = load_ids_from_file("ord_item_xm_ids.txt")[:50000]

    processor = MySQLBatchProcessor(db_host, db_port, db_user, db_password, db_name)

    try:
        processor.connect()

        # 执行步骤1
        step1_sql = """
        UPDATE CUSTOMER_LOY_RDM_ITM SET LAST_UPADTED = UTC_TIMESTAMP
        WHERE LOY_RDM_ITM_ID IN (?,?,?)
        """
        processor.execute_batch_update(step1_sql, loy_rdm_itm_ids, "步骤1 - 更新CUSTOMER_LOY_RDM_ITM")

        # # 执行步骤2
        # step2_sql = """
        # UPDATE CUSTOMER_LOY_ACRL_ITM SET LAST_UPADTED = UTC_TIMESTAMP
        # WHERE LOY_ACRL_ITM_ID IN (?,?,?)
        # """
        # processor.execute_batch_update(step2_sql, loy_acrl_itm_ids, "步骤2 - 更新CUSTOMER_LOY_ACRL_ITM")
        #
        # # 执行步骤3
        # processor.process_step3(ord_item_xm_ids)

        logger.info("所有步骤执行完成")

    except Exception as e:
        logger.error(f"主程序执行失败: {str(e)}", exc_info=True)
        sys.exit(1)
    finally:
        processor.close()


if __name__ == "__main__":
    main()